tomosphero

TomoSphero

A 3D/4D volume raytracer in spherical coordinates for arbitrary detector shape implemented in PyTorch.

Check the tutorial for instruction on using this library or examples for complete samples demonstrating forward raytracing and retrieval.

Features

  • 3D spherical raytracing with optional support for dynamic volume (4D)
  • implemented purely in PyTorch for easy integration with PyTorch's optimization and machine learning capabilities
  • support for square/circular detectors or other custom detector shapes
  • optional retrieval framework for easily defining loss functions and parametric models (currently supports only static 3D volumes)

Quickstart

pip install tomosphero
git clone https://github.com/evidlo/tomosphero && cd tomosphero
python examples/single_vantage.py


python examples/static_retrieval.py


python examples/dynamic_measurements.py

Memory Usage

This library was uses only PyTorch array operations for implementation simplicity and speed at the expense of memory consumption. The peak memory usage in GB can be approximated with examples/memory_usage.py

$ python examples/memory_usage.py

--- Parameters ---

(50, 50, 50) object
50 observations, 1 channels, (50, 100) sensor

--- Memory Usage ---

Ray coordinates memory: 4.25 GB
Object memory: 0.05 GB

Architecture

Below is a list of modules in this package and their purpose:

Forward Raytracing

  • raytracer.py - computation of voxel indices for intersecting rays, raytracing Operator
  • geometry.py - viewing geometry (detector) definitions, spherical grid definition

Retrieval

  • model.py - parameterized models for representing a density. used in retrieval
  • loss.py - some loss functions to be used in retrieval
  • retrieval.py - retrieval algorithms
  • plotting.py - functions for plotting stacks of images, retrieval losses

Running Tests

pytest tomosphero

See Also

tomosipo, which inspired parts of this library's interface.

Tutorial

This tutorial will walk through the process of setting up a tomography problem and performing a simple static reconstruction. See the examples/ directory if you want complete working scripts.

The tutorial consists of four parts:

  1. Construct a spherical grid that defines the object domain
  2. Construct view geometry which defines measurement LOS
  3. Combine grid and view geometry into an operator and take measurements
  4. Reconstruct object with iterative algorithm

Let's get started! 🮲🮳

Constructing Spherical Grids

The grid defines the physical extent and shape of the object to be raytraced or reconstructed. There are several ways to instantiate grids with the SphericalGrid class in TomoSphero:

By shape and size. e.g. for a spherical grid with 30 bins in radius/elevation/azimuth and outer radius of 10 (arbitrary) length units.

from tomosphero import SphericalGrid

grid = SphericalGrid(
    # voxels (radius, elevation, azimuth)
    shape=(30, 30, 30), 
    size_r=(0, 1),  # inner/outer grid radius
    # size_e=(0, 2*pi), # (radians, default) 
    # size_a=(-pi, pi), # (radians, default)
)
grid.plot()

An origin-centered grid with outer radius 10

Grids do not need to cover the full extent of a sphere. Below is an example of a grid defined over a small spherical wedge. This can be useful for reconstructions on a local region of a planetary atmosphere atmosphere.

from numpy import pi
from tomosphero import SphericalGrid

grid = SphericalGrid(
    # voxels (radius, elevation, azimuth)
    shape=(10, 10, 10), 
    size_r=(3, 10),  # inner/outer grid radius
    size_e=(0.4 * pi, 0.6 * pi), # (radians)
    size_a=(-.1 * pi,  .1 * pi), # (radians) 
)
grid.plot()

A grid over a spherical wedge

Alternatively, grids may be specified by defining the specific location of boundaries in radius/elevation/azimuth dimensions.

from numpy import pi
from tomosphero import SphericalGrid

grid = SphericalGrid(
    r_b=[8, 9, 10],
    e_b=[0.4*pi, 0.5*pi, 0.6*pi],
    a_b=[-0.1*pi, 0.0*pi, 0.1*pi],
)
grid.plot()

A grid defined by its boundaries

Grid shapes may be queried later with grid.shape and grid.size. See tomosphero.grid.SphericalGrid.

Constructing View Geometries

A view geometry defines the lines of sight associated with each pixel in a sensor. TomoSphero currently has 4 built-in view geometry types (contributions welcome):

ConeRectGeom is a rectilinear camera sensor, the most common sensor type. It is defined by giving sensor shape, location, orientation, and field of view:

from tomosphero import ConeRectGeom

geom = ConeRectGeom(
    shape=(64, 64), # pixels
    pos=(10, 0, 0), # sensor location
    # lookdir=(-1, 0, 0) # sensor orientation (default: towards origin)
    fov=(45, 45) # degrees
)
geom.plot()

Rectangular sensor view geometry

View geometries may be composed by addition to produce scanning orbits (e.g. circular/helical orbits), The result of this operation is a tomosphero.geometry.ViewGeomCollection object, which behaves much like a regular ViewGeom.

import numpy as np
from tomosphero import ConeRectGeom

geom = None
# sweep 360° around origin in orbit of radius 5
for w in np.linspace(0, 2*np.pi, 50):
    pos=(5*np.cos(w), 5*np.sin(w), 2)
    # combine geoms by addition
    geom += ConeRectGeom(
        pos=pos,
        shape=(100, 100), # pixels
        fov=(25, 25) # degrees
    )
anim = geom.plot()

View geometries composed into orbit

View geometry shape may be queried later with geom.shape, and its internal LOS are available through geom.ray_starts for LOS start points and geom.rays for normal vectors.

Taking Projections

The tomosphero.raytracer.Operator class is responsible for carrying out the raytracing operation, taking the grid and view geometry specified previously as arguments. At instantiation, this object will compute and store the intersection coordinates of all lines of sight with all grid boundaries and may take some time depending on the sensor and grid sizes.

TomoSphero makes extensive use of PyTorch for its autograd capabilities and GPU support. TomoSphero uses the CPU by default, but the GPU may be selected with the device kwarg.

Let's first begin by creating an Operator and defining a 3D object in a PyTorch tensor. For simplicity, let's take our object to be a sphere with a wedge missing from it:

# compute LOS intersections over grid
op = Operator(grid, geom, device='cuda')
# example phantom object - broken torus
rho = t.ones(grid.shape, device='cuda')
# sideways pacman object 
rho[:, :, 3:] = 1

After instantiation, the operator accepts PyTorch arrays with shape grid.shape (and with appropriate device) and returns raytraced measurements.

# raytrace and get measurements
y = op(rho)
# y.shape matches geom.shape

# plot measurements
from tomosphero.plotting import image_stack
anim = image_stack(y, geom)

Measurements

Reconstruction

TomoSphero is designed to be used as a component in iterative tomographic reconstruction and provides some basic building blocks for setting up an inverse problem. Given the synthetically-generated measurement $y$ from the last section, we can use this iterative approach to fit a reconstruction $\hat{\rho}$ to our measurements $y$.

Compared to more conventional techniques such as filtered back-projection, (S)ART-based methods, an iterative reconstruction approach built on an autograd framework has many benefits:

  • Works with any view geometry
    • Filtered back-projection techniques are generally limited to circular or helical orbits
  • Availability of different optimizers
    • Any PyTorch optimizer] may be used interchangeably in TomoSphero. Some optimizers may handle certain problems better than others when it comes to e.g. nonconvexity.
  • Rapid testing of different cost functions and regularizers
    • (S)ART-based methods are hard coded to a specific loss function
    • Iterative autograd techniques can work with any differentiable loss/regularizers

We begin by instantiating a tensor with requires_grad=True so PyTorch will track its gradients and set up the optimizer:

rho_hat = t.zeros(grid.shape, requires_grad=True)
optim = t.optim.Adam([rho_hat], lr=1e-1)

We will use a very basic least-squares loss function with no model parameterization or regularizers for our reconstruction:

$$\hat{\rho} = \arg \min_{\rho} ||y - F \rho||_2^2$$

where $y$ are measurements, $\rho$ is the reconstructed object, and $F$ is the tomographic operator.

for i in range(100):
    optim.zero_grad()
    # compute loss and its gradient
    loss = t.sum((y - op(rho_hat))**2)
    loss.backward()
    optim.step()

anim = image_stack(preview3d(rho))
plt.title('Ground Truth')
anim = image_stack(preview3d(rho_hat))
plt.title('Reconstruction')

We have successfully reconstructed the object! Of course, we had ideal measurement conditions with good angular coverage and no noise, but these cases can be handled by the introduction of more sophisticated models and regularization.

As mentioned above, one of the major benefits of autograd frameworks is the flexibility to change any part of the loss function or model with minimal changes to the code. Model-based retrievals provides more details on TomoSphero's tools that make experimenting with parametric models and regularizers easier.

Miscellaneous

Retrieval Framework

TomoSphero provides an optional object-oriented framework for conveniently experimenting with new loss functions and tracking loss history in iterative retrieval. The image below shows an example of a plot created with the loss tracking tools for a problem with three loss terms:

Loss plot example

Plotting Utilities

TomoSphero features extensive plotting capabilities with many objects containing a built-in .plot() function for preview. In addition, TomoSphero has a tomosphero.plotting module with several more plotting utilities:

1"""
2.. include:: ../README.md
3.. include:: ../tutorial/tutorial.md
4"""
5
6from .raytracer import Operator
7from .geometry import *
8from .model import *